package com.wangbo.trieTree;

import org.apache.commons.lang3.StringUtils;

import java.util.HashMap;

/**
 * 前缀树（HashMap实现）
 */
public class TrieTree2 {
    public static class Node {
        public int pass; // 有多少条“路径”经过该节点，“开头”和“结尾”都算“经过”
        public int end; // 有多少条“路径”以该节点结尾
        public HashMap<Integer,Node> nexts; // 该节点的后继节点

        public Node() {
            pass = 0;
            end = 0;
            nexts = new HashMap<>();
        }
    }

    public Node root;

    public TrieTree2() {
        root = new Node();
    }

    /**
     * 插入字符串 word（把字符串转换成树上的一条从root开始的“路径”）
     */
    public void insert(String word) {
        if (StringUtils.isNotBlank(word)) {
            char[] chars = word.toCharArray();
            Node node = root; // 遍历指针，从 root 遍历到 word 的最后一个字符
            node.pass++;
            int index = 0; // 字符的ASCII码
            for (int i = 0; i < chars.length; i++) {
                index = (int) chars[i];
                if (!node.nexts.containsKey(index)) {
                    node.nexts.put(index, new Node());
                }
                node = node.nexts.get(index);
                node.pass++;
            }
            node.end++;
        }
    }

    /**
     * 查询字符串 word 插入次数
     */
    public int search(String word) {
        if (StringUtils.isNotBlank(word)) {
            char[] chars = word.toCharArray();
            Node node = root; // 遍历指针，从 root 遍历到 word 的最后一个字符
            int index = 0; // 字符的ASCII码
            for (int i = 0; i < chars.length; i++) {
                index = (int) chars[i];
                if (!node.nexts.containsKey(index)) {
                    return 0;
                }
                node = node.nexts.get(index);
            }
            return node.end;
        } else {
            return 0;
        }
    }

    /**
     * 查询插入的字符串中，有多少个以指定字符串 prefix 为前缀
     */
    public int prefixNum(String prefix) {
        if (StringUtils.isNotBlank(prefix)) {
            char[] chars = prefix.toCharArray();
            Node node = root; // 遍历指针，从 root 遍历到 prefix 的最后一个字符
            int index = 0; // 字符的ASCII码
            for (int i = 0; i < chars.length; i++) {
                index = (int) chars[i];
                if (!node.nexts.containsKey(index)) {
                    return 0;
                }
                node = node.nexts.get(index);
            }
            return node.pass;
        } else {
            return 0;
        }
    }

    /**
     * 删除一个字符串 word
     */
    public void delete(String word) {
        if (search(word) > 0) {
            char[] chars = word.toCharArray();
            Node node = root; // 遍历指针，从 root 遍历到 prefix 的最后一个字符
            node.pass--;
            int index = 0; // 字符的ASCII码
            for (int i = 0; i < chars.length; i++) {
                index = (int) chars[i];
                // 如果 node 后继节点的 pass 自减后为0，说明没有其他字符串“经过”该后继节点，直接删除后继节点为根的子树，然后返回
                if (--node.nexts.get(index).pass == 0) {
                    node.nexts.remove(index);
                    return;
                }
                node = node.nexts.get(index);
            }
            node.end--;
        }
    }
}
